# -*- coding:utf-8 -*-
from ctypes import *
import pysoem
import struct
from ctypes import *
import time
import numpy as np
import mujoco
import threading


# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *


site_force = np.zeros(6)
filtered_site_force = np.zeros(6)


def xjc():
    global site_force
    global filtered_site_force
    # 创建 EtherCAT 主站，打开指定网口（请确保 enp4s0 已正确配置）
    master = pysoem.Master()
    master.open('enp4s0')

    # 扫描从站设备
    if master.config_init() == 0:
        print("未发现任何从站设备")
        master.close()
        return

    sensor = master.slaves[0]
    print("发现传感器:", sensor.name)

    # 进入 PRE-OP 状态，此状态允许通过 SDO 访问设备对象字典
    master.state = pysoem.PREOP_STATE
    master.write_state()
    # 稍作延时确保状态切换成功
    time.sleep(0.1)

    print("开始持续读取 SDO 数据，按 Ctrl+C 停止...")

    # 循环读取每个通道数据
    indices = [0x6000, 0x6001, 0x6002, 0x6003, 0x6004, 0x6005]
    channel_names = ['Fx', 'Fy', 'Fz', 'Mx', 'My', 'Mz']
    force_sensor = np.zeros(6)
    try:
        while True:
            # readings = {}
            for idx, name in zip(indices, channel_names):
                try:
                    # 读取对象 0x6000.1 ~ 0x6005.1
                    data = sensor.sdo_read(idx, 1)
                    # 根据文档，该数据类型为 REAL (32位 float)，使用小端格式解码
                    force_sensor[idx - 24576] = struct.unpack("<f", data)[0]
                    # value = struct.unpack("<f", data)[0]
                    # readings[name] = value
                    # print(f"idx {idx}, name {name}, value {value}")
                except Exception as e:
                    force_sensor[idx - 24576] = 0.0
                    print(f"Error: {e}")
                    # readings[name] = f"Error: {e}"
            # print("SDO 读取结果:", readings)
            site_force_raw = sensor_to_site(force_sensor)

            # 滤波
            filtered_site_force = alpha * site_force_raw + (1 - alpha) * filtered_site_force
            site_force = filtered_site_force.copy()
            time.sleep(0.001)  # 每 1ms 读取一次，可根据需要调整
    except KeyboardInterrupt:
        print("读取中断，程序退出")
    finally:
        master.close()


def wh_control_p():
    
    # >>>>>初始化CAN卡部分>>>
    # 设置can or canfd接口
    lcanlib = LCAN()
    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)

    # open device
    ret = lcanlib.LCAN_OpenDevice(LCAN_USBCANFD2, devIndex)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("LYS USBCANFD2 Open Device failed!")
        exit(0)
    # open success
    print("LYS USBCANFD2 Open Device Success!")


    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)

    # 初始化canfd
    ret = lcanlib.LCAN_InitCANFD(LCAN_USBCANFD2, devIndex, canIndex, canfd_init_cfg)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("LYS USBCANFD2 Init Failed!")
    else:
        print("LYS USBCANFD2 Init Success!")
    

    # 清除can通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)
    print('sleeping')
    time.sleep(0.1)
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    motor_controllers = MotorController(motor_num)

    for ID in range(motor_num):
        

        # 在线更新标志位
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_obj_init[ID], 1)
        print('sleeping')
        time.sleep(0.05)
        if ret != LCAN_STATUS_OK:
            print(f"motor {ID} init Failed!")
        else:
            print(f"motor {ID} init Success!")



        # 设置p模式
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_p_mode[ID], 1)
        print('sleeping')
        time.sleep(0.05)
        if ret != LCAN_STATUS_OK:
            print(f"Set {ID} p mode failed!")
        else:
            print(f"Set {ID} p mode success!")


    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # 控制
    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)

    # 发送与读取>>>
    # P Mode

    pvc = np.zeros((motor_num,3))
    # 每个周期，发送的目标位置和当前角度误差，最大不超过max_delta_intq


    time.sleep(3)

    start = time.time()

    ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_obj_get, motor_num)
    time.sleep(0.5)
    q_real = receive_p_state(lcanlib, devIndex, canIndex)
    qv_real = np.zeros(motor_num)

    # 读取初始点位
    model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_P.xml')) 
    data = mujoco.MjData(model)
    mocap_id = model.body("target").mocapid[0]
    mocap_pos = data.mocap_pos[mocap_id]
    mocap_quat = data.mocap_quat[mocap_id]
    
    while time.time() - start < control_time:
        loop_start = time.time()

        q_real_float = trans_intq2q(q_real)
        qv_real_float = trans_intv2v(qv_real)
        print(f"q_real_float: {q_real_float}")

        ik_1_time = time.time()
        ikq_star = ik7(q_real_float, mocap_pos, mocap_quat, qv_real_float, site_force)
        print(f"ik+admittance time cost = {time.time() - ik_1_time}")
        intikq = trans_q2intq(ikq_star) # 最终目标点（也可能是IK尚未收敛的中间变量）

        # 归零
        intikq = np.array([0,0,0,0,0,0,0])

        print(f"intikq = {intikq}")
        print(f"q_real = {q_real}")

        ctrl_intq = get_interpoint(q_real, intikq, max_dq)

        print(f"ctrl_intq = {ctrl_intq}")

        motor_controllers.set_p(ctrl_intq)
        # 发送指令
        # transfer_time1 = time.time()
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_set_p, motor_num)
        # print(f"transfer time cost: {time.time() - transfer_time1}")

        # 读取状态
        new_pvc = receive_pvc(lcanlib, devIndex, canIndex)

        if new_pvc is not None:
            pvc = new_pvc  # 仅当有有效数据时更新
            print(f"最新状态:\n {pvc}")
        else:
            print(f"get no pvc data\n")

        q_real = pvc.T[0]
        qv_real = pvc.T[1]

        elapsed_time = time.time() - loop_start
        sleep_time = control_period - elapsed_time
        if sleep_time > 0:
            time.sleep(sleep_time)
            print(f"loop time cost: {elapsed_time}")
        else:
            print(f"loop overtime: {-sleep_time}")


    # 发送与读取<<<

    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    #Close Device
    lcanlib.LCAN_CloseDevice(LCAN_USBCANFD2, devIndex)
    print("LYS USBCANFD2 Close!")



if __name__ == "__main__":
    sensor_thread = threading.Thread(target=xjc)
    sensor_thread.start()
    
    wh_control_p()
    
    sensor_thread.join()